#!/usr/bin/env python3
"""
模型文件合并脚本
将分块文件合并回原始模型文件
"""

import os
from pathlib import Path

def merge_chunks():
    """合并分块文件"""
    chunk_files = [
        "lama_fp32_part_aa",
        "lama_fp32_part_ab",
        "lama_fp32_part_ac",
        "lama_fp32_part_ad",
        "lama_fp32_part_ae",
        "lama_fp32_part_af",
        "lama_fp32_part_ag",
    ]
    
    output_file = "lama_fp32.onnx"
    
    print(f"合并 {len(chunk_files)} 个分块到 {output_file}")
    
    with open(output_file, 'wb') as outfile:
        for i, chunk_file in enumerate(chunk_files):
            print(f"处理分块 {i + 1}/{len(chunk_files)}: {chunk_file}")
            
            if not os.path.exists(chunk_file):
                print(f"错误：分块文件 {chunk_file} 不存在")
                return False
                
            with open(chunk_file, 'rb') as infile:
                while True:
                    chunk_data = infile.read(8192)
                    if not chunk_data:
                        break
                    outfile.write(chunk_data)
    
    file_size = os.path.getsize(output_file)
    print(f"合并完成！文件大小: {file_size / 1024 / 1024:.1f} MB")
    return True

if __name__ == "__main__":
    merge_chunks()
